-
Notifications
You must be signed in to change notification settings - Fork 97
feat: modify aten_bilinear from einsum to matmul #2746
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
feat: modify aten_bilinear from einsum to matmul #2746
Conversation
@microsoft-github-policy-service agree |
Note on test failureThe failure appears unrelated to the bilinear changes. The error is raised during test input generation for This suggests an RNG seed issue in the test harness rather than an op |
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #2746 +/- ##
=======================================
Coverage 70.10% 70.11%
=======================================
Files 228 228
Lines 27387 27396 +9
Branches 2785 2785
=======================================
+ Hits 19201 19210 +9
Misses 7227 7227
Partials 959 959 ☔ View full report in Codecov by Sentry. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks. @xadupre or @gramalingam for another review
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Pull request overview
This PR aims to optimize the aten_bilinear function by replacing the Einstein summation (einsum) implementation with matrix multiplication (matmul) operations for improved performance.
Key changes:
- Replaced single einsum operation with a sequence of transpose, reshape, matmul, and squeeze operations
- Added explicit shape extraction for batch dimensions and feature dimensions
- Modified computation flow to use matmul instead of einsum for the bilinear transformation
justinchuby
left a comment
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Found some issues
|
Hey @justinchuby! I fixed the batch-dim issue by switching to Shape(input1, start=0, end=-1) so all leading dims are preserved (the op requires identical batch shapes anyway). Tests still intermittently fail due to the existing RNG seed issue (same as the upsample_bilinear2d tests), which seems unrelated to this change. To better guard against regressions here, would you like me to add a test case for bilinear with batch_dim > 1? I can do that in a follow-up PR if that’s preferable. |
That would be ideal, thanks! Yes a follow up PR is preferable. |
justinchuby
left a comment
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thank you!
same as well chief |
What this does
Why
Testing
Resolves #2573